package com.onefish.toolkit.aop;

import cn.hutool.core.date.DateUtil;
import com.google.common.util.concurrent.RateLimiter;
import com.onefish.toolkit.RateLimiters;
import com.onefish.toolkit.annotation.RateLimited;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;

@Aspect
@Component
public class RateLimitAspect {
    @Pointcut("execution(* com.onefish.toolkit.controller.RateLimitController.*(..))")
    public void controllerMethods() {
    }
    //  拦截策略
    @Around("controllerMethods() && args(request,..) && @annotation(com.onefish.toolkit.annotation.RateLimited)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, HttpServletRequest request) {
        try {
            MethodSignature signature = (MethodSignature) joinPoint.getSignature();
            RateLimited rateLimited = signature.getMethod().getAnnotation(RateLimited.class);
            double permitsPerSecond = rateLimited.permitsPerSecond();
            boolean specialOption = rateLimited.specialOption();
            String uri = request.getRequestURI();
            //  具有限流策略
            if (specialOption){
                String userId = request.getParameter("userId");
                //  生成限流唯一键
                String uniKey = uri + "_" + userId;
                RateLimiter rateLimiter = RateLimiters.USER_RATE_LIMITER.computeIfAbsent(uniKey,k -> RateLimiter.create(permitsPerSecond));
                if (!rateLimiter.tryAcquire()){
                    System.out.println(DateUtil.now() + " 用户：" + userId + " 访问接口：" + uri + " 太频繁，请稍后重试，谢谢");
                    return "您访问太频繁，请稍后重试，谢谢";
                }
            } else {
                //   全局限流
                if (!RateLimiters.RATE_LIMITER.tryAcquire()){
                    System.out.println(DateUtil.now() + " 您访问接口："+ uri + " 太频繁，请稍后重试，谢谢");
                    return "您访问太频繁，请稍后重试，谢谢";
                }
            }
            System.out.println(DateUtil.now() + " 可以访问");
            return joinPoint.proceed();
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }
}
